import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import GaussianNB
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report

# -------------------------- 1. 数据加载与预处理 --------------------------
# 定义特征列名（对应 iris.names 中的 4 个数值属性）
feature_names = ['sepal_length_cm', 'sepal_width_cm', 'petal_length_cm', 'petal_width_cm']
target_name = 'class'

# 加载主数据文件 iris.data（忽略重复的 bezdekIris.data 和 Index 文件）
# 数据格式：每行是 4 个特征 + 1 个标签，用逗号分隔
df = pd.read_csv(
    'iris.data',  # 确保该文件与代码在同一目录下，否则需写绝对路径
    header=None,  # 数据无表头
    names=feature_names + [target_name],  # 手动指定列名
    encoding='utf-8'
)

# 查看数据基本信息
print("数据集形状：", df.shape)
print("\n数据集前 5 行：")
print(df.head())
print("\n数据标签分布：")
print(df[target_name].value_counts())

# 数据预处理：标签编码（字符串转数字）
le = LabelEncoder()
df[target_name] = le.fit_transform(df[target_name])  # 编码后：0=Iris-setosa, 1=Iris-versicolor, 2=Iris-virginica

# 分离特征（X）和标签（y）
X = df[feature_names].values  # 特征矩阵（150 行 × 4 列）
y = df[target_name].values    # 标签向量（150 行）

# 划分训练集和测试集（80% 训练，20% 测试，随机种子确保结果可复现）
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y  # stratify=y 保持标签分布一致
)

print(f"\n训练集大小：{X_train.shape[0]} 样本")
print(f"测试集大小：{X_test.shape[0]} 样本")

# -------------------------- 2. 训练高斯朴素贝叶斯模型 --------------------------
# 初始化高斯朴素贝叶斯分类器（适用于连续型特征）
gnb = GaussianNB()

# 用训练集训练模型
gnb.fit(X_train, y_train)

# -------------------------- 3. 模型预测与评估 --------------------------
# 用测试集预测
y_pred = gnb.predict(X_test)

# 计算核心评估指标
accuracy = accuracy_score(y_test, y_pred)  # 准确率
conf_matrix = confusion_matrix(y_test, y_pred)  # 混淆矩阵
class_report = classification_report(y_test, y_pred, target_names=le.classes_)  # 分类报告（精确率、召回率、F1）

# 输出评估结果
print("\n" + "="*50)
print("模型评估结果")
print("="*50)
print(f"准确率（Accuracy）：{accuracy:.4f}")
print("\n混淆矩阵：")
print(conf_matrix)
print("\n分类报告：")
print(class_report)

# -------------------------- 4. 可视化混淆矩阵 --------------------------
plt.figure(figsize=(8, 6))
sns.heatmap(
    conf_matrix,
    annot=True,  # 显示数值
    fmt='d',  # 数值格式为整数
    cmap='Blues',
    xticklabels=le.classes_,  # x 轴标签（原始品种名）
    yticklabels=le.classes_,  # y 轴标签（原始品种名）
    cbar_kws={'label': '样本数量'}
)
plt.xlabel('预测标签')
plt.ylabel('真实标签')
plt.title('鸢尾花分类混淆矩阵（高斯朴素贝叶斯）')
plt.tight_layout()
plt.show()

# -------------------------- 5. 预测新样本（示例） --------------------------
print("\n" + "="*50)
print("新样本预测示例")
print("="*50)
# 定义新样本（格式：[花萼长度, 花萼宽度, 花瓣长度, 花瓣宽度]）
new_samples = [
    [5.1, 3.5, 1.4, 0.2],  # 已知：Iris-setosa
    [6.4, 3.2, 4.5, 1.5],  # 已知：Iris-versicolor
    [7.2, 3.6, 6.1, 2.5]   # 已知：Iris-virginica
]

# 预测新样本
new_preds = gnb.predict(new_samples)
new_pred_names = le.inverse_transform(new_preds)  # 解码为原始品种名

# 输出预测结果
for i, (sample, pred_name) in enumerate(zip(new_samples, new_pred_names), 1):
    print(f"新样本 {i}：{sample} → 预测品种：{pred_name}")